import threading
from mitmproxy import http, websocket
import sseclient
import urllib3
import logging

def open_stream(url, headers):
    """Get a streaming response for the given event feed using urllib3."""
    http_poolman = urllib3.PoolManager()
    return http_poolman.request('GET', url, preload_content=False, headers=headers)

class SseCapture:
    def __init__(self):
        self.sse_threads = []

    def responseheaders(self,flow: http.HTTPFlow):
        pretty_url = flow.request.pretty_url
        if "eastmoney" in pretty_url:
            flow.response.stream = flow.response.headers.get('content-type', '').startswith('text/event-stream')
            # logging.info(f"responseheaders() ****NNNNN***\n pretty_url={pretty_url} \nflow.response={flow.response!r}  \n****NNNNN***\n")

    

    def request(self, flow: http.HTTPFlow):
        pretty_url = flow.request.pretty_url
        # logging.info(f"*****NNNNN request() pretty_url={pretty_url} \nrequest={flow.request!r}")
        if "sse" in pretty_url:
            # accept_header = flow.request.headers.get("accept", "")
            # logging.info(f"*****NNNNN request() accept_header={accept_header}")
            # if accept_header and "text/event-stream" in accept_header:
            logging.info(f"***NNNN request() start sse_thread pretty_url={pretty_url} \nrequest.header={flow.request.headers!r}\n request={flow.request!r}")
            # sse_thread = threading.Thread(target=self.handle_sse_request, args=(flow.request.url, dict(flow.request.headers)))
            # self.sse_threads.append(sse_thread)
            # sse_thread.start()

    def handle_sse_request(self, url, headers):
        logging.info("****NNNN handle_sse_request() .....")
        response = open_stream(url, headers)
        client = sseclient.SSEClient(response)
        stream = client.events()  # returns a generator object

        try:
            while True:
                event = next(stream)
                logging.info(f"111111111- event: {event.event} \ndata: {event.data}\nevent_id:{event.id}")
        except StopIteration:
            logging.info("2222222222- SSE stream ended")
        finally:
            response.close()  # Close the response
            logging.info("333333333- Cleaning up...")
            # Perform any other cleanup actions here

addons = [
    SseCapture()
]
